from ctypes import sizeof
import os
import numpy as np
import torch
import cv2
from torch.utils.data import Dataset

dataset_path = "dataset" # 从根目录调用此函数时path不用写

# 继承Dataset类，不用重写 len与getitem
# class ReadYolo(Dataset):
"""
:param type:输入需要读收的模型类型
    "classification":分类
    "objectdetection":目标检测
    "facedetection":人脸检测
:param phase:数据类型
    "train":训练数据
    "valid":验证数据
    "test":测试数据
:param trans:是否进行图像增强
"""


class ReadYolo(Dataset):

    def __init__(
            self,
            # type="classfication",
            phase="train",
            trans=None,
            device=None):

        super(ReadYolo, self).__init__()
        self.device = device
        self.type = type
        self.phase = phase
        self.trans = trans
        # # 有些测试数据集可能没有label
        # if phase != "test":
        #     self.label = os.listdir(os.path.join('dataset/', self.phase, 'label'))
        self.labels = os.listdir(os.path.join(dataset_path, self.phase, 'label'))
        self.imgs = os.listdir(os.path.join(dataset_path, self.phase, 'img'))
        self.img_names = list(map(lambda x: x.split('.')[0], self.imgs))

    # 重写类内函数
    def __len__(self):
        return len(self.labels)

    # list 的[] === getitem()
    def __getitem__(self, item):
        list_target = [] # 将每一行数据变成列表存入其中
        # 取出与标签对应的图片
        # 这样防止标签与图片不对应，imgs可能特别多
        # bool 对应 2 [false, true, false ,...]
        img = self.imgs[list(map(lambda x: x == self.labels[item].split('.')[0], self.img_names)).index(True)]
        img_path = os.path.join(dataset_path, self.phase, 'img', img)
        with open(os.path.join(dataset_path, self.phase, 'label', self.labels[item]), 'r') as fp:
            # 应对不规范的标注 为目标检测做准备
            for line in fp.readline():
                # 剥离换行符后行内还有数据
                if len(line.strip('\n')) > 0:
                    # 去掉首尾的空格后再split成数组
                    nums = line.strip().split(' ')
                    # 将元素浮点化后作为一个列表  *等价于list()
                    l_info = [*map(lambda x: float(x), nums)]
                    # 添加多行列表 包含类型与锚框坐标信息
                    list_target.append(l_info)

        if len(list(list_target)) == 0:
            array_target = np.array([])
        else:
            # list2nparray
            array_target = np.concatenate(list_target, axis=0).reshape(len(list_target), -1)

        pic = cv2.imread(img_path)  # array = [w h ch]

        if self.trans:
            # img,label
            pic, array_target = self.trans(pic, array_target, (224, 224)) # (h w)
            # pic [ch w h]
            # unsqueeze 扩展维度 方便训练 [b ch w h]
            return pic.unsqueeze(0).to(self.device), torch.from_numpy(array_target).to(self.device)
            # return 由于代码格式化错误第一次错在了此处
        else:
            # np.transpose(image_hwc, (2,0,1)) -->chw
            # imgchw = np.transpose(oripic, (2, 0, 1))
            # imgtorch = torch.tensor(imgchw).unsqueeze(0)
            imgchw = np.transpose(pic, (2, 0, 1))
            return torch.tensor(imgchw).unsqueeze(0), torch.from_numpy(array_target)
            # 尺寸不对应，stack expects each tensor to be equal size
            # 必须有trans


# 测试用
if __name__ == '__main__':
    from torch.utils.data import DataLoader
    from augmentation.data_augment import DataAugment
    from utils.collate import colle

    data_augment = DataAugment()
    dataset = ReadYolo(trans=data_augment)
    # dataset = ReadYolo()
    print(len(dataset))
    pic1, _ = dataset[1]
    pic2, label2 = dataset.__getitem__(1)
    # 等价
    print(pic1.shape, pic2.shape,label2)
    data = iter(DataLoader(dataset, batch_size=2, drop_last=False, collate_fn=colle))
    pics, labels = next(data)
    # print(pics, labels)
